# 手动操作船舶，使用3自由度动力学方程。
import pygame
import numpy as np
from dynamics_simulation import motionSim
import random
from typing import Dict, List, Tuple
from algo_struct.algo import rad_limit
import os
import pickle

class DrawPolygon():
    def __init__(self):
        # initial position
        self.points = [np.array([-10, -20]), np.array([10, -20]), np.array([10, 20]), np.array([0, 30]), np.array([-10, 20])]
        self.polygon = []
        # trasition
        self.rotation = 0
        self.x = 0
        self.y = 0
        self.scale=1.0

    def reset(self):
        self.points = [np.array([-10, -20]), np.array([10, -20]), np.array([10, 20]), np.array([0, 30]), np.array([-10, 20])]
        self.polygon = []
        # trasition
        self.rotation = 0
        self.x = 0
        self.y = 0
        self.scale=1.0

    def transition(self, move_x, move_y, rot_angle, scale=1.0):
        self.x += move_x
        self.y += move_y

        if self.x > SCREEN_WIDTH:
            self.x = SCREEN_WIDTH
        elif self.x < 0:
            self.x = 0
        if self.y > SCREEN_HEIGHT:
            self.y = SCREEN_HEIGHT
        elif self.y < 0:
            self.y = 0
        # print(f'self x, y is : {self.x}, {self.y}')
        self.rotation += rot_angle
        self.scale = scale
        self.polygon.clear()
        for point in self.points:
            # 先旋转
            _x = np.cos(self.rotation) * point[0] - np.sin(self.rotation) * point[1]
            _y = np.sin(self.rotation) * point[0] + np.cos(self.rotation) * point[1]
            # 后平移
            _x += self.x
            _y += self.y
            _x = self.scale * _x
            _y = self.scale * _y
            self.polygon.append([_x, _y])

    def detect_collision(self, collision: list):
        """
        检测是否发生碰撞
        :param collision: list of (x, y) - 碰撞对象的多边形顶点列表
        :return: True 如果发生碰撞, False 否则
        """
        for point in self.polygon:  # 检查当前多边形的每个点
            if is_point_inside_polygon(point, collision):
                return True

        return False
    
    def detect_overlap(self, target: list):
        # 检查self.polygon 这个list里的所有点都在凸包多边形target内部，全部都在就返回true, 有一个不在就返回false
        for point in self.polygon:  # 检查 collision 的每个点是否在当前多边形内部
            if is_point_inside_polygon(point, target) == False:
                return False
        # 所有点都在目标多边形内部 返回True
        return True
    
def is_point_inside_polygon(point, polygon):
    """
    使用射线法判断点是否在多边形内部
    :param point: (x, y) - 点的坐标
    :param polygon: list of (x, y) - 多边形顶点的顺序列表
    :return: True 如果点在多边形内部，否则 False
    """
    x, y = point
    count = 0
    n = len(polygon)
    
    for i in range(n):
        x1, y1 = polygon[i]
        x2, y2 = polygon[(i + 1) % n]
        
        # 检查点是否在多边形的边界上
        if (y1 <= y < y2 or y2 <= y < y1) and x <= max(x1, x2):
            xinters = (y - y1) * (x2 - x1) / (y2 - y1 + 1e-10) + x1
            if abs(x - xinters) < 1e-10:  # 在边上
                return True
            if x < xinters:  # 射线穿过
                count += 1
    
    return count % 2 == 1


def draw_text(surface, text, pos, font_size=24, color=(0, 0, 0)):
    """在屏幕上绘制文本"""
    font = pygame.font.Font(None, font_size)
    text_surface = font.render(text, True, color)
    surface.blit(text_surface, pos)

def draw_info_panel(surface, observation,exp_speed, exp_angle):
    """绘制信息面板显示船舶状态"""
    panel_x = 10
    panel_y = 10
    line_height = 25
    
    # 背景面板
    panel_width = 300
    panel_height = 300
    panel_surface = pygame.Surface((panel_width, panel_height))
    panel_surface.set_alpha(200)
    panel_surface.fill((255, 255, 255))
    surface.blit(panel_surface, (panel_x, panel_y))
    
    # 显示各种状态信息
    info_texts = [
        f"Boat State:",
        f"Position X: {observation.get('x', 0):.2f} m",
        f"Position Y: {observation.get('y', 0):.2f} m", 
        f"Position Z: {observation.get('z', 0):.2f} m", 
        f"X Speed: {observation.get('u', 0):.2f} m/s",
        f"Y Speed: {observation.get('v', 0):.2f} m/s",
        f"Z Speed: {observation.get('w', 0):.2f} m/s",
    ]
    
    for i, text in enumerate(info_texts):
        color = (0, 0, 0) if i > 0 else (0, 0, 255)  # 标题用蓝色
        font_size = 20 if i > 0 else 24
        draw_text(surface, text, (panel_x + 5, panel_y + 5 + i * line_height), font_size, color)

def draw_controls_help(surface):
    """绘制控制说明"""
    help_x = SCREEN_WIDTH - 250
    help_y = 10
    
    # 背景面板
    panel_surface = pygame.Surface((240, 120))
    panel_surface.set_alpha(200)
    panel_surface.fill((200, 255, 200))
    surface.blit(panel_surface, (help_x, help_y))
    
    help_texts = [
        "Description:",
        "WERTYU: Add voltage",# six voltage
        "SDFGHJ: Decrease voltage",
        "SPACE: PAUSE/CONTINUE",
        "R: RESET"
    ]
    
    for i, text in enumerate(help_texts):
        color = (0, 100, 0) if i > 0 else (0, 0, 100)
        font_size = 18 if i > 0 else 20
        draw_text(surface, text, (help_x + 5, help_y + 5 + i * 20), font_size, color)

# pygame setup
pygame.init()
pygame.font.init()  # 初始化字体模块
SCREEN_WIDTH = 1000
SCREEN_HEIGHT = 600
screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
pygame.display.set_caption("USV Simulation")
clock = pygame.time.Clock()
running = True
dt = 0

# Predefined some colors
BLUE = (0, 0, 255)
RED = (255, 0, 0)
GREEN = (0, 255, 0)
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
YELLOW = (255, 255, 0)

RECT_WIDTH = 50
RECT_HEIGHT = 100

# 绘制障碍物
enemy_pos = DrawPolygon()
enemy_pos.transition(SCREEN_WIDTH * 2/5, SCREEN_HEIGHT * 2/5, np.pi/2)

# 绘制玩家船舶
player_pos = DrawPolygon()
player_pos.transition(SCREEN_WIDTH/2, SCREEN_HEIGHT/2, np.pi)

# 目标区域
goal_pos = DrawPolygon()
goal_pos.transition(SCREEN_WIDTH * 2/5, SCREEN_HEIGHT * 3/5, np.pi/2, scale=1.5)

# USV dynamic 
# 非线性kt参数
K=3.1429
T=18.348
deltam=6.8207e-2
alpha=6.8595e-3

delta_h = 0.1

# boat运动仿真类输入参数
my_mmg = motionSim.BoatMotionSim(delta_h=delta_h, name='rov')
x = np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0], dtype=np.float32)# 初始ROV位置
u = np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0], dtype=np.float32)# 初始ROV速度
voltage = np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0], dtype=np.float32)# 初始控制电压信号

observation=dict()
memory_demo = list()
count =0
exp_speed = random.choice([-0.5, -0.2, 0, -0.2, 0.5])# 指定list中随机选择
exp_angle = random.randint(-90, 90)
epoch = 1

paused = False
show_trail = True
trail_points = []  # 存储轨迹点

while running:
    # poll for events
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        elif event.type == pygame.KEYDOWN:
            if event.key == pygame.K_SPACE:
                paused = not paused
            elif event.key == pygame.K_r:
                # 重置状态
                x = np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0], dtype=np.float32)
                u = np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0], dtype=np.float32)
                voltage = np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0], dtype=np.float32)
                
                player_pos.reset()
                player_pos.transition(SCREEN_WIDTH/2, SCREEN_HEIGHT/2, np.pi)
                trail_points.clear()
                exp_speed = random.choice([-0.5, -0.2, 0, -0.2, 0.5])
                exp_angle = random.randint(-90, 90)
            elif event.key == pygame.K_t:
                show_trail = not show_trail
                
    if not paused:
        # 手动驱动
        keys = pygame.key.get_pressed()
        if keys[pygame.K_w]:
            voltage[0] += 0.1
        if keys[pygame.K_e]:
            voltage[1] += 0.1
        if keys[pygame.K_r]:
            voltage[2] += 0.1 
        if keys[pygame.K_t]:
            voltage[3] += 0.1
        if keys[pygame.K_y]:
            voltage[4] += 0.1
        if keys[pygame.K_u]:
            voltage[5] += 0.1
        if keys[pygame.K_s]:
            voltage[0] -= 0.1
        if keys[pygame.K_d]:
            voltage[1] -= 0.1
        if keys[pygame.K_f]:
            voltage[2] -= 0.1
        if keys[pygame.K_g]:
            voltage[3] -= 0.1
        if keys[pygame.K_h]:
            voltage[4] -= 0.1
        if keys[pygame.K_j]:
            voltage[5] -= 0.1

        # 限制控制量范围
        voltage = np.clip(voltage, -1.0, 1.0)
        
        print(f'voltage is {voltage}')
        
        # 计算更新 位置 速度  m/s 和 rad
        x, u, nu, acc = my_mmg.run_ROV(x=x, u=u, v=voltage)
        x[3] = rad_limit(x[3])
        x[4] = rad_limit(x[4])
        x[5] = rad_limit(x[5]) # 限制到-pi到pi
        
        # 更新观测值
        observation = {
            'x': x[0],
            'y': x[1],
            'z': x[2],
            'phi': x[3] * 180.0/np.pi,
            'theta': x[4] * 180.0/np.pi, # rad to deg
            'psi': x[5] * 180.0/np.pi,
            'u': u[0],
            'v': u[1],
            'w': u[2],
            'voltage': voltage,
        }
        
        # 更新船舶位置
        player_pos.transition(nu[1]*dt, -nu[0]* dt, nu[5]*dt)
        # 记录轨迹点
        if len(trail_points) == 0 or np.linalg.norm([player_pos.x - trail_points[-1][0], player_pos.y - trail_points[-1][1]]) > 5:
            trail_points.append((player_pos.x, player_pos.y))
            # 限制轨迹点数量，避免内存过多占用
            if len(trail_points) > 500:
                trail_points.pop(0)

    # 绘制场景
    screen.fill(WHITE)
    
    # 绘制轨迹
    if show_trail and len(trail_points) > 1:
        for i in range(1, len(trail_points)):
            alpha = i / len(trail_points)  # 轨迹渐变效果
            color = (int(255 * (1-alpha)), int(255 * alpha), 0)  # 从红到黄的渐变
            pygame.draw.line(screen, color, trail_points[i-1], trail_points[i], 2)
    
    # 绘制玩家船舶
    pygame.draw.polygon(screen, RED, player_pos.polygon)
    
    # 绘制船舶中心点
    pygame.draw.circle(screen, BLUE, (int(player_pos.x), int(player_pos.y)), 3)
    
    # 绘制船舶朝向指示线
    heading_length = 40
    heading_end_x = player_pos.x + heading_length * np.cos(player_pos.rotation)
    heading_end_y = player_pos.y + heading_length * np.sin(player_pos.rotation)
    pygame.draw.line(screen, BLUE, (player_pos.x, player_pos.y), (heading_end_x, heading_end_y), 3)
    
    # 绘制速度矢量
    if 'u' in observation and 'v' in observation:
        speed_scale = 20  # 速度显示比例
        vel_end_x = player_pos.x + observation['u'] * speed_scale
        vel_end_y = player_pos.y - observation['v'] * speed_scale  # pygame y轴向下
        pygame.draw.line(screen, GREEN, (player_pos.x, player_pos.y), (vel_end_x, vel_end_y), 2)
        # 绘制速度矢量箭头
        if abs(observation['u']) > 0.1 or abs(observation['v']) > 0.1:
            pygame.draw.circle(screen, GREEN, (int(vel_end_x), int(vel_end_y)), 4)
    
    # 绘制期望航向指示
    exp_heading_length = 60
    exp_heading_rad = np.radians(exp_angle)
    exp_end_x = player_pos.x + exp_heading_length * np.cos(exp_heading_rad)
    exp_end_y = player_pos.y + exp_heading_length * np.sin(exp_heading_rad)
    pygame.draw.line(screen, YELLOW, (player_pos.x, player_pos.y), (exp_end_x, exp_end_y), 2)
    
    # 绘制信息面板
    if 'x' in observation:
        draw_info_panel(screen, observation, exp_speed, exp_angle)
    
    # 绘制控制说明
    draw_controls_help(screen)
    
    # 绘制暂停状态
    if paused:
        pause_surface = pygame.Surface((SCREEN_WIDTH, SCREEN_HEIGHT))
        pause_surface.set_alpha(100)
        pause_surface.fill((0, 0, 0))
        screen.blit(pause_surface, (0, 0))
        draw_text(screen, "PAUSED - Press SPACE to continue", 
                 (SCREEN_WIDTH//2 - 150, SCREEN_HEIGHT//2), 32, (255, 255, 255))
    
    # 绘制状态栏
    status_y = SCREEN_HEIGHT - 30
    status_texts = [
        f"Time: {count*dt:.1f}s",
        f"Epoch: {epoch}",
        f"Trail: {'ON' if show_trail else 'OFF'} (T to toggle)"
    ]
    
    for i, text in enumerate(status_texts):
        draw_text(screen, text, (10 + i * 150, status_y), 18, (50, 50, 50))
    
    # 更新显示
    pygame.display.flip()

    # 控制帧率和时间步长
    dt = 0.1
    clock.tick(60)  # 60 FPS
    count += 1
    
pygame.quit()
